Skip to content

cuda-gen: Add CUDA Graph capture and replay for composite operators - #1987

Open
Nafees01 wants to merge 18 commits into
CEED:mainfrom
Nafees01:cuda-graph-dev
Open

cuda-gen: Add CUDA Graph capture and replay for composite operators#1987
Nafees01 wants to merge 18 commits into
CEED:mainfrom
Nafees01:cuda-graph-dev

Conversation

@Nafees01

@Nafees01 Nafees01 commented Jul 8, 2026

Copy link
Copy Markdown

Summary

Adds CUDA Graph capture and replay to the 'cuda-gen' backend for composite operators. On the second apply, the kernel sequence is captured into a cudaGraph; all subsequent applies replay the instantiated graph. QFunction contexts are re-synced to device before each replay so time- and load-dependent parameters stay correct.

Sub-operators that cannot be captured, for example, contact operators that allocate memory (cudaMalloc) during apply, which CUDA doesn't allow inside a capture. When this happens, the capture is properly closed to leave the stream in a clean state, any leftover CUDA error is cleared, and that operator falls back to running normally on every call.

Validated against both /gpu/cuda/ref and cuda-gen without graphs across static, quasistatic, contact, multi-material, and dynamic elasticity examples. Strain energy and displacements match to ~1e-13 in every case, with zero graph fallbacks on capturable operators.

Performance-wise, graphs are within ±1% of cuda-gen without graphs across all sizes, so no regression, but no real speedup either. The reason is that cuda-gen already fuses everything into a few large kernels, so there's very little launch overhead left for graphs to remove.

Validated correctness against /gpu/cuda/ref and cuda-gen (CEED_DISABLE_GRAPH=1) on the following Ratel examples:

  • ex01-static-elasticity-linear-mms
  • ex02-quasistatic-elasticity-linear-platen (contact)
  • ex02-quasistatic-elasticity-neo-hookean-current-sinker (multi-material)
  • ex03-dynamic-elasticity-mooney-rivlin-current
    Strain energy matches to ~1e-13 in all cases. Zero graph fallbacks on capturable operators.

@zatkins-dev

Copy link
Copy Markdown
Collaborator

Sub-operators that cannot be captured, for example, contact operators that allocate memory (cudaMalloc) during apply

Are there specific operators you were unable to use in Ratel? None of the QFunctions there should alloc during operator application.

@zatkins-dev

Copy link
Copy Markdown
Collaborator

There's some style errors here, be sure to make format each commit

@zatkins-dev zatkins-dev left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm sure you're still working on this, but here's an initial review to correct some high-level design issues and removed functionality that I think should be fixed.

Comment thread backends/cuda-gen/ceed-cuda-gen-operator.c
Comment thread backends/cuda-gen/ceed-cuda-gen-operator.c Outdated
Comment thread backends/cuda-gen/ceed-cuda-gen-operator.c Outdated
Comment thread backends/cuda-gen/ceed-cuda-gen-operator.c Outdated
Comment thread backends/cuda-gen/ceed-cuda-gen-operator.c Outdated
Comment thread backends/cuda-gen/ceed-cuda-gen-operator.c Outdated
Comment thread backends/cuda-gen/ceed-cuda-gen-operator.c Outdated
Comment thread backends/cuda-gen/ceed-cuda-gen-operator.c Outdated
Comment thread backends/cuda-gen/ceed-cuda-gen-operator.c Outdated
Comment thread backends/cuda-gen/ceed-cuda-gen-operator.c Outdated
@zatkins-dev

Copy link
Copy Markdown
Collaborator

I would be interested to see if this speeds up e.g. /gpu/cuda/ref, when recording the graph over the entire CeedOperatorApplyCore_Cuda_Ref function

@jeremylt

Copy link
Copy Markdown
Member

Note - merging main into your branch instead of rebasing means I'd recommend we squash-merge this PR

@Nafees01

Copy link
Copy Markdown
Author

Sub-operators that cannot be captured, for example, contact operators that allocate memory (cudaMalloc) during apply

Are there specific operators you were unable to use in Ratel? None of the QFunctions there should alloc during operator application.

The only example where the contact (Nitsche) sub-operator couldn't be captured was the quasistatic contact case (ex02-quasistatic-elasticity-linear-platen). Everything else I tried static, dynamic, multi-material, neo-Hookean, Mooney-Rivlin captured and replayed fine.

And you're right that the QFunction doesn't allocate; the cudaMalloc comes from libCEED's own device-side handling of the contact data, not the QFunction itself. It runs either way correctly; the contact operators just fall back to a normal apply rather than being replayed.

@zatkins-dev

Copy link
Copy Markdown
Collaborator

Sub-operators that cannot be captured, for example, contact operators that allocate memory (cudaMalloc) during apply

Are there specific operators you were unable to use in Ratel? None of the QFunctions there should alloc during operator application.

The only example where the contact (Nitsche) sub-operator couldn't be captured was the quasistatic contact case (ex02-quasistatic-elasticity-linear-platen). Everything else I tried static, dynamic, multi-material, neo-Hookean, Mooney-Rivlin captured and replayed fine.

And you're right that the QFunction doesn't allocate; the cudaMalloc comes from libCEED's own device-side handling of the contact data, not the QFunction itself. It runs either way correctly; the contact operators just fall back to a normal apply rather than being replayed.

This is a pretty big red flag to me. The Nitsche contact operators don't actually do anything out of the ordinary on the libCEED side. Them failing to work means that something about this setup (very possibly something with passive inputs and outputs) is violating the libCEED interface assumptions.

@Nafees01

Copy link
Copy Markdown
Author

Sub-operators that cannot be captured, for example, contact operators that allocate memory (cudaMalloc) during apply

Are there specific operators you were unable to use in Ratel? None of the QFunctions there should alloc during operator application.

The only example where the contact (Nitsche) sub-operator couldn't be captured was the quasistatic contact case (ex02-quasistatic-elasticity-linear-platen). Everything else I tried static, dynamic, multi-material, neo-Hookean, Mooney-Rivlin captured and replayed fine.
And you're right that the QFunction doesn't allocate; the cudaMalloc comes from libCEED's own device-side handling of the contact data, not the QFunction itself. It runs either way correctly; the contact operators just fall back to a normal apply rather than being replayed.

This is a pretty big red flag to me. The Nitsche contact operators don't actually do anything out of the ordinary on the libCEED side. Them failing to work means that something about this setup (very possibly something with passive inputs and outputs) is violating the libCEED interface assumptions.

Thanks for flagging this. On re-checking linear-platen with CEED_DEBUG=1, the contact operators appear to fail earlier than graph capture: cuda-gen hits an NVRTC compile error (GradTransposeTensor3dFlattened defined twice in cuda-shared-basis-tensor-flattened-templates.h) and then falls back to /gpu/cuda/ref. So this looks like a pre-existing cuda-gen JIT issue triggered by the contact operator's basis configuration, rather than something specific to Nitsche or CUDA graphs.

@zatkins-dev

Copy link
Copy Markdown
Collaborator

On re-checking linear-platen with CEED_DEBUG=1, the contact operators appear to fail earlier than graph capture: cuda-gen hits an NVRTC compile error (GradTransposeTensor3dFlattened defined twice in cuda-shared-basis-tensor-flattened-templates.h) and then falls back to /gpu/cuda/ref

Sweet, good catch. I'll fix that separately.

Comment thread backends/cuda-gen/ceed-cuda-gen-operator.c Outdated
Comment thread backends/cuda-gen/ceed-cuda-gen-operator.c
Comment thread backends/cuda-gen/ceed-cuda-gen-operator.c Outdated
Comment thread backends/cuda-gen/ceed-cuda-gen-operator.c
Comment thread backends/cuda-ref/ceed-cuda-ref-qfunctioncontext.c
Comment thread backends/cuda-ref/ceed-cuda-ref-vector.c
Comment thread backends/cuda-gen/ceed-cuda-gen-operator.c Outdated
Comment thread interface/ceed-cuda.c
Comment thread interface/ceed-cuda.c
Comment on lines 15 to +18
@brief Set CUDA function pointer to evaluate action at quadrature points

If the backend does not support `CUfunction` pointers for QFunctions, the call succeeds without effect.
When unsupported, a message is emitted via `CeedDebug`.

@jeremylt jeremylt Jul 28, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
@brief Set CUDA function pointer to evaluate action at quadrature points
If the backend does not support `CUfunction` pointers for QFunctions, the call succeeds without effect.
When unsupported, a message is emitted via `CeedDebug`.
@brief Set CUDA function pointer to evaluate action at quadrature points.
If the backend does not support `CUfunction` pointers for QFunctions, then the call succeeds without effect.
When unsupported, a message is emitted via `CeedDebug`.

minor

Comment thread interface/ceed-cuda.c
Comment on lines +37 to +40
@brief Enable or disable CUDA Graph capture/replay for a `CeedOperator`

If the backend does not support CUDA Graphs for operators, the call succeeds without effect.
When unsupported, a message is emitted via `CeedDebug`.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
@brief Enable or disable CUDA Graph capture/replay for a `CeedOperator`
If the backend does not support CUDA Graphs for operators, the call succeeds without effect.
When unsupported, a message is emitted via `CeedDebug`.
@brief Enable or disable CUDA Graph capture/replay for a `CeedOperator`.
If the backend does not support CUDA Graphs for operators, then the call succeeds without effect.
When unsupported, a message is emitted via `CeedDebug`.

@jeremylt
jeremylt requested a review from zatkins-dev July 28, 2026 06:08
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants